import os
import gc
import ctypes
import logging
from typing import Any, Dict, List, Union

# 配置日志
logger = logging.getLogger(__name__)

class SecureMemory:
    """安全内存管理类，提供内存敏感数据的安全擦除功能"""
    
    @staticmethod
    def secure_erase(data: Any) -> Any:
        """
        安全擦除数据，根据数据类型选择合适的擦除方法
        
        Args:
            data: 要擦除的数据
            
        Returns:
            擦除后的数据（通常为空或零值）
        """
        if data is None:
            return None
            
        if isinstance(data, bytearray):
            # 直接覆写字节数组
            for i in range(len(data)):
                data[i] = 0
            return data
            
        elif isinstance(data, bytes):
            # 对于不可变的bytes类型，返回相同长度的零字节
            return b'\x00' * len(data)
            
        elif isinstance(data, str):
            # 对于字符串，返回空字符串
            return ""
            
        elif isinstance(data, (list, tuple)):
            # 递归处理列表或元组中的每个元素
            return type(data)([SecureMemory.secure_erase(item) for item in data])
            
        elif isinstance(data, dict):
            # 递归处理字典中的每个值
            return {k: SecureMemory.secure_erase(v) for k, v in data.items()}
            
        elif isinstance(data, (int, float)):
            # 对于数字，返回0
            return 0
            
        else:
            # 对于其他类型，尝试设为None
            try:
                return None
            except:
                logger.warning(f"无法安全擦除类型: {type(data)}")
                return data
    
    @staticmethod
    def secure_clear_dict(data_dict: Dict) -> Dict:
        """
        安全清除字典中的所有值
        
        Args:
            data_dict: 要清除的字典
            
        Returns:
            清除后的字典（相同的键，值被擦除）
        """
        if not isinstance(data_dict, dict):
            return {}
            
        cleared_dict = {}
        for key, value in data_dict.items():
            cleared_dict[key] = SecureMemory.secure_erase(value)
            
        return cleared_dict
    
    @staticmethod
    def secure_clear_object(obj: Any) -> None:
        """
        尝试安全清除对象的敏感属性
        
        Args:
            obj: 要清除属性的对象
        """
        if obj is None:
            return
            
        try:
            # 获取对象的所有属性
            for attr_name in dir(obj):
                # 跳过特殊属性和方法
                if attr_name.startswith('__') or callable(getattr(obj, attr_name)):
                    continue
                    
                try:
                    attr_value = getattr(obj, attr_name)
                    # 对于可变类型的属性，进行安全擦除
                    if isinstance(attr_value, (str, bytes, bytearray, list, dict, set)):
                        setattr(obj, attr_name, SecureMemory.secure_erase(attr_value))
                except AttributeError:
                    # 跳过只读属性
                    pass
        except Exception as e:
            logger.error(f"清除对象属性时出错: {str(e)}")
    
    @staticmethod
    def force_gc() -> None:
        """强制进行垃圾回收，释放内存"""
        try:
            gc.collect()
            logger.debug("强制垃圾回收完成")
        except Exception as e:
            logger.error(f"强制垃圾回收失败: {str(e)}")
    
    @staticmethod
    def secure_exit(session_state: Any = None) -> None:
        """
        执行安全退出操作
        
        Args:
            session_state: 要清理的会话状态对象
        """
        try:
            # 清理会话状态
            if session_state is not None:
                SecureMemory.secure_clear_object(session_state)
                
            # 强制垃圾回收
            SecureMemory.force_gc()
            
            logger.info("已安全清理内存")
        except Exception as e:
            logger.error(f"安全退出错误: {str(e)}") 